from __future__ import absolute_import

# import loopdev,mount
#import tempfile
#import shutil
#import time
import os

#from os import makedirs
from os.path import join, dirname #, isdir
from subprocess import check_call, PIPE

from forest.api import get_tool
#from forest.tools.system.loopdev import LoopDev
from forest.tools.system.mount import Mount
from forest.tools.system.copytree import copytree
from forest.registry import register_tool
from forest.tools.system.mksparse import mksparse
from forest.tools.image.loops import find_unused_loop_device

######## Functions for making filesystem images ##########################

def do_grafts(grafts, dest, preserve = True):
    '''Copy each of the items listed in grafts into dest.
    If the key ends with '/' it's assumed to be a directory which should be
    created, otherwise just the leading directories will be created.'''
    for imgpath, filename in grafts.items():
        if imgpath[-1] == '/':
            targetdir = join(dest, imgpath)
            imgpath = imgpath[:-1]
        else:
            targetdir = join(dest, dirname(imgpath))
        if not os.path.isdir(targetdir):
            os.makedirs(targetdir)
        if os.path.isdir(filename):
            copytree(filename, join(dest, imgpath), preserve)
        else:
            get_tool('system.copy')(filename, join(dest, imgpath))

@register_tool('image.round_to_blocks')
def round_to_blocks(size, blocksize):
    '''If size isn't a multiple of blocksize, round up to the next multiple'''
    diff = size % blocksize
    if diff or not size:
        size += blocksize - diff
    return size

# TODO: move filesystem data outside this function
def estimate_size(rootdir, graft = None, fstype = None, blocksize = 4096, overhead = 128):

    if graft is None:
        graft = {}

    getsize = lambda f: os.lstat(f).st_size
    if fstype == "btrfs":
        overhead = 64*1024 # don't worry, it's all sparse
    if fstype in ("vfat", "msdos"):
        blocksize = 2048
        getsize = lambda f: os.stat(f).st_size # no symlinks, count as copies
    total = overhead*blocksize
    dirlist = graft.values()
    if rootdir:
        dirlist.append(rootdir)
    for root in dirlist:
        for top, dirs, files in os.walk(root):
            for f in files + dirs:
                total += round_to_blocks(getsize(join(top, f)), blocksize)
    if fstype == "btrfs":
        total = max(256*1024*1024, total) # btrfs minimum size: 256MB
    return total



@register_tool('image.mkfsimage')
def mkfsimage(fstype, rootdir, outfile, size = None, mkfsargs = None, graft = None): #mountargs = "",
    ''' Generic filesystem image creation function.

        fstype should be a filesystem type - "mkfs.${fstype}" must exist.
        graft should be a dict: {"some/path/in/image": "local/file/or/dir"};

        if the path ends with a '/' it's assumed to be a directory.
        Will raise CalledProcessError if something goes wrong.
    '''

    if mkfsargs is None:
        mkfsargs = []

    #if mountargs is None:
    #    mountargs = ""

    if graft is None:
        graft = {}

    preserve = (fstype not in ("msdos", "vfat"))
    if not size:
        size = estimate_size(rootdir, graft, fstype)

    #tmpf = tempfile.NamedTemporaryFile()

    mksparse(outfile, size)

    l = find_unused_loop_device()
    l.set_target(outfile)
    l.set_size(size)

    with l as loopdev:
        print "mounted %s" % outfile, loopdev

        check_call(["mkfs.%s" % fstype] + mkfsargs + [loopdev], stdout = PIPE, stderr = PIPE)
        check_call(['sync'])

    with Mount(outfile, "loop") as mnt:
        print "after Mount"

        if rootdir:
            print "rootdir exists1"
            copytree(rootdir, mnt, preserve)

        do_grafts(graft, mnt, preserve)
    print "after exiting LoopDev"

# convenience functions with useful defaults

#mountarg_default = "shortname=winnt,umask=0077"
@register_tool('image.mkdosimage')
def mkdosimg(rootdir, outfile, size = None, label = "", graft = None): #, mountargs = mountarg_default

    mkfsimage("msdos", rootdir, outfile, size , #mountargs = mountargs,
              mkfsargs = ["-n", label], graft = graft)

@register_tool('image.mkext4image')
def mkext4img(rootdir, outfile, size = None, label = "", graft = None):


    mkfsimage("ext3", rootdir, outfile, size , #mountargs = mountargs,
              mkfsargs = ["-L", label, "-b", "1024", "-m", "0"], graft = graft)

@register_tool('image.btrfs')
def mkbtrfsimg(rootdir, outfile, size = None, label = "", graft = None):

    mkfsimage("btrfs", rootdir, outfile, size, #mountargs = mountargs,
               mkfsargs = ["-L", label], graft = graft)

def mkhfsimg(rootdir, outfile, size = None, label = "", graft = None):

    mkfsimage("hfs", rootdir, outfile, size, #mountargs = mountargs,
              mkfsargs = ["-v", label], graft = graft)
